# from misc.classifier.word_dict import word_dict
# from scipy.sparse import csr_matrix
# from sklearn.feature_extraction import DictVectorizer
from sklearn.feature_extraction.text import CountVectorizer
 # train_test_split
# from sklearn.linear_model import LogisticRegression
# from sklearn.model_selection import train_test_split
# from sklearn.model_selection import StratifiedKFold
# from sklearn.model_selection import GridSearchCV
from sklearn.feature_extraction.text import TfidfTransformer
import pickle
import sys


if __name__ == '__main__':
    data_file = sys.argv[1]
    v_file = sys.argv[2]
    tfidf_file = sys.argv[3]
    out_file = sys.argv[4]

    target = []
    X = []
    with open(data_file) as fd:
        for l in fd:
            d = l.strip().split('\t')
            if len(d) == 3:
                t, x1, x2 = d
                target.append(t.replace('__label__', ''))
                X.append(x1 + ' ' + x2)
            if len(d) == 2:
                t, x = d
                target.append(t.replace('__label__', ''))
                X.append(x)

    # 将词转化为词向量
    CV = CountVectorizer(token_pattern=r'\b\w+\b', min_df=1)
    X = CV.fit_transform(X)
    # print(X[0:10])
    TFIDF = TfidfTransformer()
    X = TFIDF.fit_transform(X)

    with open(v_file, 'wb+') as fd:
        pickle.dump(CV, fd)

    with open(tfidf_file, 'wb+') as fd:
        pickle.dump(TFIDF, fd)

    # 输出特征数据集
    with open(out_file, 'wb+') as fd:
        pickle.dump((X, target), fd)